import os

from dotenv import load_dotenv
from openai import OpenAI
from langchain.embeddings.base import Embeddings
from typing import List

from langchain_openai import ChatOpenAI

# 加载.env文件
load_dotenv()

# 创建模型
model = ChatOpenAI(
	model='qwen-plus',
	api_key=os.getenv("OPENAI_API_KEY"),
	base_url=os.getenv("BASE_URL"),
)


def get_model():
	return model


class AliEmbeddings(Embeddings):
	def __init__(self):
		self.client = OpenAI(
			api_key=os.getenv("OPENAI_API_KEY"),
			base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
		)

	def embed_documents(self, texts: List[str]) -> List[List[float]]:
		response = self.client.embeddings.create(
			model="text-embedding-v4",
			input=texts,
		)
		return [item.embedding for item in response.data]

	def embed_query(self, text: str) -> List[float]:
		response = self.client.embeddings.create(
			model="text-embedding-v4",
			input=text,
		)
		return response.data[0].embedding
